package com.lagou.niancheng.core.transaction;

import com.lagou.niancheng.core.exceptions.LagouMybatisException;

import javax.sql.DataSource;
import java.io.Closeable;
import java.io.IOException;
import java.sql.Connection;
import java.sql.SQLException;

/**
 * 默认事务管理器
 *
 * @author 大犬
 * @date 2020/3/27
 */
public class DefaultTransactionManager implements TransactionManager {

    private DataSource dataSource;
    private ThreadLocal<CloseableConnection> connectionHolder;
    private Boolean autoCommit;

    public DefaultTransactionManager(DataSource dataSource,
                                     Boolean autoCommit) {
        this.dataSource = dataSource;
        this.connectionHolder = new ThreadLocal<>();
        this.autoCommit = autoCommit;
    }

    @Override
    public Connection openTransaction() {
        try {
            Connection connection = dataSource.getConnection();
            connection.setAutoCommit(autoCommit);
            CloseableConnection closeableConnection = new CloseableConnection(connectionHolder, connection);
            connectionHolder.set(closeableConnection);
            return connection;
        } catch (SQLException e) {
            throw new LagouMybatisException("open transactional connection failed", e);
        }
    }

    @Override
    public Connection getConnection() {
        CloseableConnection closeableConnection = connectionHolder.get();
        if (closeableConnection != null) {
            return closeableConnection.connection;
        }
        try {
            return dataSource.getConnection();
        } catch (SQLException e) {
            throw new LagouMybatisException("open connection failed", e);
        }
    }

    @Override
    public void commit() {
        try (CloseableConnection closeableConnection = connectionHolder.get()) {
            if (closeableConnection != null) {
                Connection connection = closeableConnection.connection;
                if (!connection.getAutoCommit()) {
                    connection.commit();
                }
            }
        } catch (Exception e) {
            throw new LagouMybatisException("transaction commit failed", e);
        }
    }

    @Override
    public void rollback() {
        try (CloseableConnection closeableConnection = connectionHolder.get()) {
            if (closeableConnection != null) {
                Connection connection = closeableConnection.connection;
                if (!connection.getAutoCommit()) {
                    connection.rollback();
                }
            }
        } catch (Exception e) {
            throw new LagouMybatisException("transaction rollback failed", e);
        }
    }


    private static class CloseableConnection implements Closeable {

        private ThreadLocal<CloseableConnection> connectionHolder;
        private Connection connection;

        public CloseableConnection(ThreadLocal<CloseableConnection> connectionHolder,
                                   Connection connection) {
            this.connectionHolder = connectionHolder;
            this.connection = connection;
        }

        @Override
        public void close() throws IOException {
            connectionHolder.remove();
        }
    }
}
